package p.ithorns.framework.common.codec;

import p.ithorns.framework.common.model.Pair;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA加解密工具类
 */
public class RSACodec {

    // 算法名
    private final static String ALGORITHM = "AES";

    private static final Base64.Encoder ENCODER = Base64.getEncoder();

    private static final Base64.Decoder DECODER = Base64.getDecoder();

    /**
     * RSA最大加密明文大小
     */
    private static final int MAX_ENCRYPT_BLOCK = 117;

    /**
     * RSA最大解密密文大小
     */
    private static final int MAX_DECRYPT_BLOCK = 128;

    /**
     * 字符集
     */
    private static String ENCODING = "UTF-8";

    /**
     * 获取密钥对
     *
     * @return 密钥对
     */
    public static Map<String, String> getPublicPrivateKeyMap() throws Exception {
        KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
        generator.initialize(1024);
        // 获取原始密钥对
        KeyPair keyPair = generator.generateKeyPair();
        // 获取公钥、私钥
        String privateKey = ENCODER.encodeToString(keyPair.getPrivate().getEncoded());
        String publicKey = ENCODER.encodeToString(keyPair.getPublic().getEncoded());
        // 封装密钥对
        Map<String, String> keyMap = new HashMap<>();
        keyMap.put("privateKey", privateKey);
        keyMap.put("publicKey", publicKey);
        return keyMap;
    }

    public static String encodeKey(Key key) {
        return ENCODER.encodeToString(key.getEncoded());
    }

    public static Pair<String, String> getKeyPair() throws Exception {
        KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
        generator.initialize(1024);
        // 获取原始密钥对
        KeyPair keyPair = generator.generateKeyPair();
        // 获取公钥、私钥
        String privateKey = new String(ENCODER.encode(keyPair.getPrivate().getEncoded()));
        String publicKey = new String(ENCODER.encode(keyPair.getPublic().getEncoded()));
        // 封装密钥对
        return Pair.of(privateKey, publicKey);
    }

    /**
     * 获取私钥
     *
     * @param privateKey 私钥字符串
     * @return 私钥
     */
    public static PrivateKey getPrivateKey(String privateKey) throws Exception {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        byte[] decodedKey = DECODER.decode(privateKey.getBytes());
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodedKey);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 获取公钥
     *
     * @param publicKey 公钥字符串
     * @return 公钥
     */
    public static PublicKey getPublicKey(String publicKey) throws Exception {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        byte[] decodedKey = DECODER.decode(publicKey.getBytes());
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodedKey);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * RSA 私钥加密
     *
     * @param data       待加密数据
     * @param privateKey 私钥
     * @return 加密数据
     */
    public static String encryptPrivate(String data, String privateKey) throws Exception {
        return encrypt(data, getPrivateKey(privateKey));
    }

    /**
     * RSA 公钥加密
     *
     * @param data      待加密数据
     * @param publicKey 公钥
     * @return 加密数据
     */
    public static String encryptPublic(String data, String publicKey) throws Exception {
        return encrypt(data, getPublicKey(publicKey));
    }

    /**
     * RSA 加密
     *
     * @param data -待加密数据
     * @param key  -秘钥
     * @return String 加密结果
     * @throws Exception -
     */
    public static String encrypt(String data, Key key) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        int inputLen = data.getBytes().length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offset = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段加密
        while (inputLen - offset > 0) {
            if (inputLen - offset > MAX_ENCRYPT_BLOCK) {
                cache = cipher.doFinal(data.getBytes(), offset, MAX_ENCRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(data.getBytes(), offset, inputLen - offset);
            }
            out.write(cache, 0, cache.length);
            i++;
            offset = i * MAX_ENCRYPT_BLOCK;
        }
        byte[] encryptedData = out.toByteArray();
        out.close();
        // 获取加密内容使用base64进行编码,并以UTF-8为标准转化成字符串
        // 加密后的字符串
        return ENCODER.encodeToString(encryptedData);
    }

    /**
     * RSA 公钥加密-私钥解密
     *
     * @param data       待解密数据
     * @param privateKey 私钥
     * @return 解密数据
     */
    public static String decrypt(String data, String privateKey) throws Exception {
        return decrypt(data, getPrivateKey(privateKey));
    }

    /**
     * RSA 私钥加密-公钥解密
     *
     * @param data      待解密数据
     * @param publicKey 公钥
     * @return 解密数据
     */
    public static String decryptPublic(String data, String publicKey) throws Exception {
        return decrypt(data, getPublicKey(publicKey));
    }

    /**
     * RSA 解密
     *
     * @param data -待解密数据
     * @param key  -秘钥
     * @return String 解密结果
     * @throws Exception -
     */
    public static String decrypt(String data, Key key) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, key);
        byte[] dataBytes = DECODER.decode(data);
        int inputLen = dataBytes.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offset = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段解密
        while (inputLen - offset > 0) {
            if (inputLen - offset > MAX_DECRYPT_BLOCK) {
                cache = cipher.doFinal(dataBytes, offset, MAX_DECRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(dataBytes, offset, inputLen - offset);
            }
            out.write(cache, 0, cache.length);
            i++;
            offset = i * MAX_DECRYPT_BLOCK;
        }
        out.close();
        // 解密后的内容
        return out.toString(ENCODING);
    }

    /**
     * 私钥签名
     *
     * @param data          -待签名数据
     * @param privateKeyStr -私钥
     * @return String 签名
     * @throws Exception -
     */
    public static String sign(String data, String privateKeyStr) throws Exception {
        PrivateKey privateKey = getPrivateKey(privateKeyStr);
        return sign(data, privateKey);
    }

    /**
     * 签名
     *
     * @param data       待签名数据
     * @param privateKey 私钥
     * @return 签名
     */
    public static String sign(String data, PrivateKey privateKey) throws Exception {
        byte[] keyBytes = privateKey.getEncoded();
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        PrivateKey key = keyFactory.generatePrivate(keySpec);
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initSign(key);
        signature.update(data.getBytes());
        return ENCODER.encodeToString(signature.sign());
    }

    /**
     * 签名验证
     *
     * @param data         -待验证数据
     * @param publicKeyStr -公钥
     * @param sign         -签名
     * @return boolean 是否验签通过
     * @throws Exception -
     */
    public static boolean verify(String data, String publicKeyStr, String sign) throws Exception {
        PublicKey publicKey = getPublicKey(publicKeyStr);
        return verify(data, publicKey, sign);
    }

    /**
     * 验签
     *
     * @param srcData   原始字符串
     * @param publicKey 公钥
     * @param sign      签名
     * @return 是否验签通过
     * @throws Exception -
     */
    public static boolean verify(String srcData, PublicKey publicKey, String sign) throws Exception {
        byte[] keyBytes = publicKey.getEncoded();
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        PublicKey key = keyFactory.generatePublic(keySpec);
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initVerify(key);
        signature.update(srcData.getBytes());
        return signature.verify(DECODER.decode(sign.getBytes()));
    }

}